package com.cloud.mall.util;

import io.netty.buffer.ByteBufAllocator;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.CharBuffer;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicReference;

/**
 * @Author: 梁二东
 * @Date: 2024/12/8 10:38
 * @Description:
 **/
public class GatewayBodyUtils {

    /**
     * 获取请求中的body
     *
     * @param req
     * @return
     */
    public static String getBody(ServerHttpRequest req) {
        if (checkBodyHeader(req)) {
            AtomicReference<String> requestBody = new AtomicReference<>("");
            RecorderServerHttpRequestDecorator requestDecorator = new RecorderServerHttpRequestDecorator(req);
            Flux<DataBuffer> body = requestDecorator.getBody();
            body.subscribe(buffer -> {
                CharBuffer charBuffer = StandardCharsets.UTF_8.decode(buffer.asByteBuffer());
                requestBody.set(charBuffer.toString());
            });
            return requestBody.get();
        }
        return null;
    }

    /**
     * 修改设置body
     *
     * @param body     不能为空
     * @param exchange
     * @param chain
     * @return
     */
    public static Mono<Void> setBody(String body, ServerWebExchange exchange, GatewayFilterChain chain) {
        DataBuffer bodyDataBuffer = stringBuffer(body);
        Flux<DataBuffer> bodyFlux = Flux.just(bodyDataBuffer);
        MediaType contentType = exchange.getRequest().getHeaders().getContentType();
        ServerHttpRequest mutatedRequest = new ServerHttpRequestDecorator(exchange.getRequest()) {
            @Override
            public HttpHeaders getHeaders() {
                HttpHeaders httpHeaders = new HttpHeaders();
                int length = body.getBytes().length;
                httpHeaders.putAll(super.getHeaders());
                httpHeaders.remove(HttpHeaders.CONTENT_TYPE);
                httpHeaders.remove(HttpHeaders.CONTENT_LENGTH);
                httpHeaders.setContentLength(length);
                httpHeaders.set(HttpHeaders.CONTENT_TYPE, contentType.toString());
                // 设置CONTENT_TYPE
                return httpHeaders;
            }

            @Override
            public Flux<DataBuffer> getBody() {
                return bodyFlux;
            }
        };
        return chain.filter(exchange.mutate().request(mutatedRequest).build());
    }


    /**
     * 校验body头
     *
     * @param req
     * @return
     */
    public static boolean checkBodyHeader(ServerHttpRequest req) {
        HttpHeaders headers = req.getHeaders();
        MediaType contentType = headers.getContentType();
        long contentLength = headers.getContentLength();
        if (contentLength > 0 && !MediaType.MULTIPART_FORM_DATA_VALUE.startsWith(contentType.toString())) {
            if (MediaType.APPLICATION_JSON.equals(contentType) || MediaType.APPLICATION_JSON_UTF8.equals(contentType) || contentType.toString().contains(MediaType.APPLICATION_FORM_URLENCODED_VALUE)) {
                return true;
            }
        }
        return false;
    }

    public static DataBuffer stringBuffer(String value) {
        byte[] bytes = value.getBytes(StandardCharsets.UTF_8);
        NettyDataBufferFactory nettyDataBufferFactory = new NettyDataBufferFactory(ByteBufAllocator.DEFAULT);
        DataBuffer buffer = nettyDataBufferFactory.allocateBuffer(bytes.length);
        buffer.write(bytes);
        return buffer;
    }

    public static class RecorderServerHttpRequestDecorator extends ServerHttpRequestDecorator {
        private final List<DataBuffer> dataBuffers = new ArrayList<>();

        public RecorderServerHttpRequestDecorator(ServerHttpRequest delegate) {
            super(delegate);
            super.getBody().map(dataBuffer -> {
                dataBuffers.add(dataBuffer);
                return dataBuffer;
            }).subscribe();
        }

        @Override
        public Flux<DataBuffer> getBody() {
            return copy();
        }

        private Flux<DataBuffer> copy() {
            return Flux.fromIterable(dataBuffers)
                    .map(buf -> buf.factory().wrap(buf.asByteBuffer()));
        }
    }


}
